#!/usr/bin/env python3
import asyncio
import websockets
import json
import sys
import paho.mqtt.client as mqtt
from typing import Dict, Any
import time
import logging

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("smart-classroom-mcp-client")

class SmartClassroomMCPClient:
    def __init__(self):
        self.sensor_data = {
            'temperature': 25.0,
            'humidity': 60.0,
            'light': '正常',
            'updated_at': time.time()
        }
        
        # MQTT 配置
        self.mqtt_config = {
            'host': '39.96.163.93',
            'port': 8083,
            'device_id': 'anran'
        }
        
        # 小智 MCP 配置
        self.xiaozhi_endpoint = "wss://api.xiaozhi.me/mcp/?token=eyJhbGciOiJFUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VySWQiOjU5MTQ0MSwiYWdlbnRJZCI6OTEzMDA1LCJlbmRwb2ludElkIjoiYWdlbnRfOTEzMDA1IiwicHVycG9zZSI6Im1jcC1lbmRwb2ludCIsImlhdCI6MTc2MTk5NzEzOCwiZXhwIjoxNzkzNTU0NzM4fQ.m6QT5uMFJA0E1v2T87YYmYUaGkZAm1Np33h0OMsQMVtak3Fmt0Q0dg2iA-w2txvPyl1jX0xT6qI05ljIqRWWrg"
        
        self.mqtt_client = None
        self.websocket = None
        self.is_connected = False
        self.setup_mqtt()
    
    def setup_mqtt(self):
        """初始化 MQTT 客户端"""
        try:
            logger.info("正在初始化 MQTT 连接...")
            # 使用更新的回调API版本
            self.mqtt_client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2, transport="websockets")
            self.mqtt_client.on_connect = self.on_mqtt_connect
            self.mqtt_client.on_message = self.on_mqtt_message
            self.mqtt_client.on_disconnect = self.on_mqtt_disconnect
            
            # 设置连接超时
            self.mqtt_client.connect_async(self.mqtt_config['host'], self.mqtt_config['port'])
            self.mqtt_client.loop_start()
            
            logger.info("MQTT 客户端初始化完成")
            
        except Exception as e:
            logger.error(f"MQTT 初始化失败: {e}")
    
    def on_mqtt_connect(self, client, userdata, flags, reason_code, properties):
        """MQTT 连接回调"""
        if reason_code == 0:
            logger.info("MQTT 连接成功!")
            client.subscribe(f"{self.mqtt_config['device_id']}/sensor/+")
            client.subscribe(f"{self.mqtt_config['device_id']}/state/+")
        else:
            logger.error(f"MQTT 连接失败，错误码: {reason_code}")
    
    def on_mqtt_disconnect(self, client, userdata, disconnect_flags, reason_code, properties):
        """MQTT 断开连接回调"""
        logger.warning(f"MQTT 连接断开，错误码: {reason_code}")
    
    def on_mqtt_message(self, client, userdata, message):
        """MQTT 消息回调"""
        try:
            topic = message.topic
            payload = message.payload.decode()
            logger.info(f"收到 MQTT 消息: {topic} -> {payload}")
            
            if topic == f"{self.mqtt_config['device_id']}/sensor/dht11":
                parts = payload.split('_')
                if len(parts) == 2:
                    self.sensor_data['temperature'] = float(parts[0])
                    self.sensor_data['humidity'] = float(parts[1])
                    self.sensor_data['updated_at'] = time.time()
            elif topic == f"{self.mqtt_config['device_id']}/sensor/light":
                self.sensor_data['light'] = '较强' if payload == '1' else '较弱'
                self.sensor_data['updated_at'] = time.time()
                
        except Exception as e:
            logger.error(f"处理 MQTT 消息错误: {e}")
    
    async def connect_to_xiaozhi(self):
        """连接到小智 MCP 服务"""
        try:
            logger.info("正在连接到小智 MCP 服务...")
            self.websocket = await websockets.connect(
                self.xiaozhi_endpoint,
                ping_interval=20,
                ping_timeout=10,
                close_timeout=10
            )
            
            # 首先等待服务器发送的初始化请求
            init_request = await self.websocket.recv()
            logger.info(f"收到小智初始化请求: {init_request}")
            
            init_data = json.loads(init_request)
            
            # 响应初始化请求
            init_response = {
                "jsonrpc": "2.0",
                "id": init_data.get("id", 1),
                "result": {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {
                        "tools": {"listChanged": False},
                        "resources": {"listChanged": False},
                        "logging": {},
                        "prompts": {}
                    },
                    "serverInfo": {
                        "name": "smart-classroom-mcp",
                        "version": "1.0.0"
                    }
                }
            }
            
            await self.websocket.send(json.dumps(init_response))
            logger.info("✅ 初始化响应已发送")
            
            # 等待 initialized 通知
            initialized_msg = await self.websocket.recv()
            logger.info(f"收到初始化完成通知: {initialized_msg}")
            
            logger.info("✅ 成功连接到小智 MCP 服务")
            self.is_connected = True
            return True
            
        except Exception as e:
            logger.error(f"连接小智 MCP 服务失败: {e}")
            self.is_connected = False
            return False
    
    def get_tools_list(self):
        """获取工具列表"""
        return {
            "tools": [
                {
                    "name": "control_light",
                    "description": "控制教室灯光开关 - 开灯或关灯",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "action": {
                                "type": "string",
                                "enum": ["on", "off"],
                                "description": "开灯(on)或关灯(off)"
                            }
                        },
                        "required": ["action"]
                    }
                },
                {
                    "name": "control_beep", 
                    "description": "控制报警器开关 - 开启或关闭报警",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "action": {
                                "type": "string", 
                                "enum": ["on", "off"],
                                "description": "开启报警(on)或关闭报警(off)"
                            }
                        },
                        "required": ["action"]
                    }
                },
                {
                    "name": "get_sensor_data",
                    "description": "获取智慧教室传感器数据 - 温度、湿度、光照",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "sensor_type": {
                                "type": "string",
                                "enum": ["all", "temperature", "humidity", "light"],
                                "description": "要查询的传感器类型: all(全部), temperature(温度), humidity(湿度), light(光照)"
                            }
                        }
                    }
                },
                {
                    "name": "get_system_status",
                    "description": "获取系统状态和连接信息",
                    "inputSchema": {
                        "type": "object",
                        "properties": {}
                    }
                }
            ]
        }
    
    async def handle_xiaozhi_requests(self):
        """处理来自小智 MCP 服务的请求"""
        try:
            async for message in self.websocket:
                try:
                    request = json.loads(message)
                    logger.info(f"收到小智 MCP 请求: {request}")
                    
                    response = await self.process_xiaozhi_request(request)
                    if response:
                        await self.websocket.send(json.dumps(response))
                        
                except json.JSONDecodeError as e:
                    logger.error(f"JSON 解析错误: {e}")
                except Exception as e:
                    logger.error(f"处理请求错误: {e}")
                    
        except websockets.exceptions.ConnectionClosed as e:
            logger.warning(f"小智 MCP 连接已关闭: {e}")
            self.is_connected = False
        except Exception as e:
            logger.error(f"处理请求时发生错误: {e}")
            self.is_connected = False
    
    async def process_xiaozhi_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """处理小智 MCP 请求"""
        method = request.get("method")
        params = request.get("params", {})
        request_id = request.get("id", 1)
        
        logger.info(f"处理小智 MCP 请求: {method}")
        
        try:
            if method == "tools/call":
                return await self.handle_tool_call(request_id, params)
            elif method == "tools/list":
                return await self.handle_tools_list(request_id)
            elif method == "ping":
                return {"jsonrpc": "2.0", "id": request_id, "result": {}}
            elif method == "notifications/initialized":
                # 忽略初始化通知，我们已经处理了
                return None
            else:
                logger.warning(f"未知方法: {method}")
                return self.error_response(request_id, "MethodNotFound", f"未知方法: {method}")
                
        except Exception as e:
            logger.error(f"处理小智请求错误: {e}")
            return self.error_response(request_id, "InternalError", str(e))
    
    async def handle_tools_list(self, request_id: int) -> Dict:
        """处理工具列表请求"""
        logger.info("处理工具列表请求")
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": self.get_tools_list()
        }
    
    async def handle_tool_call(self, request_id: int, params: Dict) -> Dict:
        """处理工具调用请求"""
        name = params.get("name")
        arguments = params.get("arguments", {})
        
        logger.info(f"小智调用工具: {name} 参数: {arguments}")
        
        try:
            if name == "control_light":
                result = self.control_light(arguments)
            elif name == "control_beep":
                result = self.control_beep(arguments)
            elif name == "get_sensor_data":
                result = self.get_sensor_data(arguments)
            elif name == "get_system_status":
                result = self.get_system_status()
            else:
                return self.error_response(request_id, "ToolNotFound", f"未知工具: {name}")
            
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "result": {
                    "content": [
                        {
                            "type": "text",
                            "text": result
                        }
                    ]
                }
            }
        except Exception as e:
            return self.error_response(request_id, "ToolExecutionFailed", str(e))
    
    def error_response(self, request_id: int, code: str, message: str) -> Dict:
        """错误响应"""
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "error": {
                "code": code,
                "message": message
            }
        }
    
    # 工具实现方法
    def control_light(self, args: Dict) -> str:
        """控制灯光"""
        action = args.get('action')
        if action not in ['on', 'off']:
            raise ValueError("action 必须是 'on' 或 'off'")
        
        message = 'a' if action == 'on' else 'b'
        if self.mqtt_client and self.mqtt_client.is_connected():
            self.mqtt_client.publish(f"{self.mqtt_config['device_id']}/cmd", message)
            return f"✅ 灯光控制命令已发送: {'开启灯光' if action == 'on' else '关闭灯光'}"
        else:
            return f"⚠️ MQTT 未连接，模拟操作: {'开启灯光' if action == 'on' else '关闭灯光'}"
    
    def control_beep(self, args: Dict) -> str:
        """控制报警器"""
        action = args.get('action')
        if action not in ['on', 'off']:
            raise ValueError("action 必须是 'on' 或 'off'")
        
        message = 'c' if action == 'on' else 'd'
        if self.mqtt_client and self.mqtt_client.is_connected():
            self.mqtt_client.publish(f"{self.mqtt_config['device_id']}/cmd", message)
            return f"✅ 报警器控制命令已发送: {'开启报警' if action == 'on' else '关闭报警'}"
        else:
            return f"⚠️ MQTT 未连接，模拟操作: {'开启报警' if action == 'on' else '关闭报警'}"
    
    def get_sensor_data(self, args: Dict) -> str:
        """获取传感器数据"""
        sensor_type = args.get('sensor_type', 'all')
        
        if sensor_type == 'all':
            data = self.sensor_data
        else:
            data = {sensor_type: self.sensor_data.get(sensor_type)}
        
        sensor_info = f"📊 传感器数据 ({sensor_type}):\n"
        for key, value in data.items():
            if value is not None and key != 'updated_at':
                if key == 'temperature':
                    sensor_info += f"🌡️ 温度: {value}°C\n"
                elif key == 'humidity':
                    sensor_info += f"💧 湿度: {value}%RH\n"
                elif key == 'light':
                    sensor_info += f"💡 光照: {value}\n"
        
        sensor_info += f"\n🕒 最后更新: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(self.sensor_data['updated_at']))}"
        
        return sensor_info
    
    def get_system_status(self) -> str:
        """获取系统状态"""
        status_info = "🔧 系统状态:\n"
        status_info += f"🔗 MQTT 连接: {'✅ 已连接' if self.mqtt_client and self.mqtt_client.is_connected() else '❌ 未连接'}\n"
        status_info += f"🤖 小智 MCP: {'✅ 已连接' if self.is_connected else '❌ 未连接'}\n"
        status_info += f"🏠 MQTT 服务器: {self.mqtt_config['host']}\n"
        status_info += f"📱 设备ID: {self.mqtt_config['device_id']}\n"
        
        # 传感器状态
        active_sensors = [k for k, v in self.sensor_data.items() if v is not None and k != 'updated_at']
        status_info += f"📊 活跃传感器: {', '.join(active_sensors) if active_sensors else '无'}\n"
        
        return status_info
    
    async def run(self):
        """运行主循环"""
        logger.info("🚀 启动智能教室 MCP 客户端...")
        
        while True:
            try:
                # 连接到小智 MCP 服务
                connected = await self.connect_to_xiaozhi()
                if connected:
                    # 启动处理任务
                    await self.handle_xiaozhi_requests()
                else:
                    logger.warning("连接失败，10秒后重试...")
                    await asyncio.sleep(10)
                    
            except Exception as e:
                logger.error(f"主循环错误: {e}")
                await asyncio.sleep(10)
            finally:
                # 清理资源
                if self.websocket:
                    await self.websocket.close()
                    self.websocket = None
                self.is_connected = False

async def main():
    """主函数"""
    client = SmartClassroomMCPClient()
    await client.run()

if __name__ == "__main__":
    # 安装必要的依赖
    try:
        import websockets
        import paho.mqtt.client as mqtt
    except ImportError as e:
        print(f"请安装必要的依赖: pip install websockets paho-mqtt")
        sys.exit(1)
    
    asyncio.run(main())